package gold.contest;

import gold.utils.InputUtil;

import java.util.*;

/**
 * Created by fanzhenyu02 on 2021/12/10.
 * common problem solver template.
 */
public class LC6017 {
    public long startExecuteTime = System.currentTimeMillis();


    /*
     * @param 此题目参考了别人代码
     * 不是因为复杂，而是因为经典
     * 未来需要再次复习此道题目
     * 等差数列求和我是没想到的
     * @return:
     */
    class Solution {
        public long minimalKSum(int[] nums, int k) {
            TreeSet<Integer> treeSet = new TreeSet<>();
            for (int num : nums) {
                treeSet.add(num);
            }
            long ans = 0;
            int cur = 1, next = 0;
            for (int num : treeSet) {
                next = cur + k - 1;
                if (next < num) {
                    ans += (long) k * (cur + next) / 2;
                    k = 0;
                    break;
                } else {
                    next = num - 1;
                    if (next >= cur) {
                        int x = next - cur + 1;
                        ans += (long) x * (cur + next) / 2;
                        k -= x;
                    }
                    cur = num + 1;
                }
            }
            if (k > 0) {
                next = cur + k - 1;
                ans += (long) k * (cur + next) / 2;
            }
            return ans;
        }
    }


    class Solution_Slow_Thinking {
        public long minimalKSum(int[] nums, int k) {
            Arrays.sort(nums);
            long sum = 0;
            int curInsert = 1, i = 0;
            while (k-- > 0) {
                if (i < nums.length) {
                    if (curInsert == nums[i]) {
                        i++;
                        while (i < nums.length && (nums[i] == nums[i - 1] || nums[i] == nums[i - 1] + 1)) i++;
                        curInsert = nums[i - 1] + 1;
                    }
                    sum += curInsert++;
                } else {
                    // i 已经到达尾部
                    sum += curInsert++;
                }
            }

            return sum;
        }
    }

    public void run() {
        Solution solution = new Solution();
        System.out.println(solution.minimalKSum(InputUtil.toIntegerArray("[1,4,25,10,25]"), 2));
        System.out.println(solution.minimalKSum(InputUtil.toIntegerArray("[5,6]"), 6));
        System.out.println(solution.minimalKSum(InputUtil.toIntegerArray("[53,41,90,33,84,26,50,32,63,47,66,43,29,88,71,28,83]"), 76));
    }

    public static void main(String[] args) throws Exception {
        LC6017 an = new LC6017();
        an.run();

        System.out.println("\ncurrent solution total execute time: " + (System.currentTimeMillis() - an.startExecuteTime) + " ms.");
    }
}
